package org.kiwi.data.adapter.druid;

import com.alibaba.druid.pool.DruidConnectionHolder;
import com.alibaba.druid.pool.DruidPooledConnection;
import org.kiwi.data.TableShardingStrategyHolder;
import org.kiwi.data.vo.TableShardingParameter;

import java.sql.CallableStatement;
import java.sql.PreparedStatement;
import java.sql.SQLException;

import static org.kiwi.data.util.ShardingUtil.convertSql;

/**
 * Created by jack08.liu on 2016/12/22.
 */
public class DruidPooledConnectionAdapter extends DruidPooledConnection {

    public DruidPooledConnectionAdapter(DruidConnectionHolder holder) {
        super(holder);
    }

    @Override
    public PreparedStatement prepareStatement(String sql) throws SQLException {
        return super.prepareStatement(convertSqlIfNecessary(sql));
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency)
            throws SQLException {
        return super.prepareStatement(convertSqlIfNecessary(sql), resultSetType, resultSetConcurrency);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency,
                                              int resultSetHoldability) throws SQLException {
        return super.prepareStatement(convertSqlIfNecessary(sql), resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
        return super.prepareStatement(convertSqlIfNecessary(sql), columnIndexes);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
        return super.prepareStatement(convertSqlIfNecessary(sql), columnNames);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
        return super.prepareStatement(convertSqlIfNecessary(sql), autoGeneratedKeys);
    }

    @Override
    public CallableStatement prepareCall(String sql) throws SQLException {
        return super.prepareCall(convertSqlIfNecessary(sql));
    }

    @Override
    public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency,
                                         int resultSetHoldability) throws SQLException {
        return super.prepareCall(convertSqlIfNecessary(sql), resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
        return super.prepareCall(convertSqlIfNecessary(sql), resultSetType, resultSetConcurrency);
    }

    private String convertSqlIfNecessary(String rawSql) {
        TableShardingParameter tableShardingParameter = TableShardingStrategyHolder.getTableShardingParameter();
        if (tableShardingParameter != null) {
            return convertSql(rawSql, tableShardingParameter);
        }
        return rawSql;
    }

}
